/*
 * Alloy Analyzer
 * Copyright (c) 2002 Massachusetts Institute of Technology
 *
 * The Alloy Analyzer is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * The Alloy Analyzer is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with the Alloy Analyzer; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 */

package alloy.transform;

import java.util.Collections;
import java.util.Iterator;
import java.util.Set;
import java.util.TreeSet;

import alloy.ast.ASTDepthFirstVisitor;
import alloy.ast.ArrowMultExpr;
import alloy.ast.BinaryFormula;
import alloy.ast.CompOp;
import alloy.ast.ComprehensionExpr;
import alloy.ast.Decl;
import alloy.ast.Decls;
import alloy.ast.ElemFormula;
import alloy.ast.Expr;
import alloy.ast.Formula;
import alloy.ast.FormulaSeq;
import alloy.ast.Formulas;
import alloy.ast.IfThenElseIntExpr;
import alloy.ast.IntExpr;
import alloy.ast.LiteralIntExpr;
import alloy.ast.LogicOp;
import alloy.ast.Multiplicity;
import alloy.ast.MultiplicityExpr;
import alloy.ast.MultiplicityFormula;
import alloy.ast.Node;
import alloy.ast.QuantifiedFormula;
import alloy.ast.Quantifier;
import alloy.ast.SetMultExpr;
import alloy.ast.SumIntExpr;
import alloy.ast.Variable;
import alloy.ast.VariableExpr;
import alloy.type.UnionType;
import alloy.util.Dbg;

/**
 * <p>A visitor which desugars the multiplicity constraints of
 * declarations appearing in various contexts.
 * <p>Assumes that only some and all quantified formulas remain.
 */

public class DesugarDeclsVisitor extends ASTDepthFirstVisitor {

    public void visit(QuantifiedFormula quantFormula_) {
        // first, desugar decls in children
        visit((Node)quantFormula_);

	//System.out.println("BEFORE BODY DESUGARING");
	//System.out.println(quantFormula_.nodeString());
	Formula body = quantFormula_.getFormula();
	FormulaSeq declConstraints = null;
	Formula replacement = null;

	Quantifier quant = quantFormula_.getQuantifier();
	if (quant.equals(Quantifier.ALL)) {
	    // desugar formulas of the form 
	    // all x1, ..., xn : lone S | F(x1,...,xn) to formulas of the form
	    // all x1, ..., xn : one S  | F(x1,...,xn) && F(none,...,xn) && 
	    //                            F(x1,none,x3,...,xn) && ... && F(none,none,...,none)
	    // (see _getLoneConstraints)
	    Formula desugarAllLone = _getLoneConstraints(quantFormula_.getDecls(), 
							 body, LogicOp.AND);

	    // desugar modified decls ("lone" became "one")
	    declConstraints = getDeclConstraints(quantFormula_.getDecls());

	    // in a universally quantified formula, the truth of the declaration
	    // constraints implies the truth of the body
	    replacement = (!declConstraints.getFormulas().isLeaf() ? 
			   declConstraints.implies(desugarAllLone) : 
			   desugarAllLone);
	}
	else if (quant.equals(Quantifier.SOME)) {
	    /* NOTE BY EMINA:  this desugaring is no longer used because it is broken.
	     * We simply do not desugar 'lone' when inside an existential quantifier --
	     * we let the skolemizer handle it if it can */
	    // desugar formulas of the form 
	    // some x1, ..., xn : lone S | F(x1,...,xn) to formulas of the form
	    // some x1, ..., xn : one S  | F(x1,...,xn) || F(none,...,xn) ||
	    //                             F(x1,none,x3,...,xn) || ... || F(none,none,...,none)
	    // (see _getLoneConstraints)
	    /*
	    Formula desugarSomeLone = _getLoneConstraints(quantFormula_.getDecls(),
							  body, LogicOp.OR);
        */
        Formula desugarSomeLone = body;
	    // desugar modified decls ("lone" became "one")
	    declConstraints = getDeclConstraints(quantFormula_.getDecls());

	    // in an existentially quantified formula, the declaration
	    // constraints are conjoined with the body
	    replacement = (!declConstraints.getFormulas().isLeaf() ? 
			   declConstraints.and(desugarSomeLone) : 
			   desugarSomeLone);
	}
	else {
	    Dbg.fatal("Quantifiers other than ALL and SOME should have been desugared.");
	}

	declConstraints.annotate("implicit constraints from declarations");
	replacement.annotateReplacing("", body);
	quantFormula_.setFormula(replacement);
	
	
	//System.out.println("AFTER BODY DESUGARING");
	//System.out.println(quantFormula_.nodeString());

    }

    /**
     * For all decl of the form x1, ... , xn : lone S in decls, produces the formula
     * f(x1,...,xn) op f(none,...,xn) op f(x1,none,x3,...,xn) op ... op f(none,none,...,none)
     * where we generate all possible combinations of inputs to f such that any of
     * its arguments may be replaced by none.
     * For example, if decls contains the declaration x1, x2 : lone S, we'll get
     * f(x1,x2) op f(none,x2) op f(x1,none) op f(none,none)
     **/
    
    private static Formula _getLoneConstraints(Decls decls, Formula f, LogicOp op) {
	Formula ret = f;
        for (Iterator declIter = decls.getDeclIter(); declIter.hasNext(); ) {

	    Decl decl = (Decl) declIter.next();

	
	    MultiplicityExpr multExpr = (MultiplicityExpr) decl.getExpr();
	    
	    if (multExpr.getMult().equals(Multiplicity.LONE)) {
		
		
		// iterate through the variables; at each step,
		// copy the entire accumulated expression, and 
		// replace the current variable, in the copy, with none.
		// then, create a new formula, ret op (ret with var replaced)
		Expr none = UnionType.empty(multExpr.getType().arity()).toExpr();
		
		for (Iterator varIter = decl.getVariables().getVariableIter();
		     varIter.hasNext();) {
		    Variable var = (Variable)varIter.next();
		    
		    ExprAndTypeReplaceVisitor replacer = new
			ExprAndTypeReplaceVisitor(Collections.singletonMap(var.nodeString(),
									   none),
						  var.nodeString(), "none");
		    
		    Formula retCopy = (Formula) ((Formula) ret.copy()).applyReturnVisitor(replacer);
		    
		    // if no replacements were performed, we don't do anything
		    if (replacer.replacementMade()) {
			ret = new BinaryFormula(ret, op, retCopy);
		    }
		}

		// change the multiplicity to "one" instead of "lone"
		multExpr.setMult(Multiplicity.ONE);
	    }
	
	}

	return ret;
    }
    

    public void visit(ComprehensionExpr compExpr_) {
        // first, desugar decls in children
	visit((Node)compExpr_);
        FormulaSeq declConstraints = getDeclConstraints(compExpr_.getDecls());
        if (!declConstraints.getFormulas().isLeaf()) {
            declConstraints.annotate("implicit constraints from declarations");
            // in a comprehension expression, the constraints
            // are conjoined to the body formula
            Formula body = compExpr_.getFormula();
            Formula replacement = declConstraints.and(body);
            replacement.annotateReplacing("", body);
            compExpr_.setFormula(replacement);            
        }
    }

    public void visit(SumIntExpr sumIntExpr_) {
        // first, desugar decls in children
	visit((Node)sumIntExpr_);
        FormulaSeq declConstraints = getDeclConstraints(sumIntExpr_.getDecls());
        if (!declConstraints.getFormulas().isLeaf()) {
            declConstraints.annotate("implicit constraints from declarations");
            // in a sumintexpr, constraints are added as follows:
            // sum decls | if constraints then body else 0
            IntExpr body = sumIntExpr_.getExpr();
            IntExpr replacement = new IfThenElseIntExpr(declConstraints,
                                                        (IntExpr)body,
                                                  new LiteralIntExpr(0));
            replacement.annotateReplacing("", body);
            sumIntExpr_.setExpr(replacement);
        }
    }
    
    public static FormulaSeq getDeclConstraints(Decls decls_) {
        Formulas allConstraints = new Formulas();
        Iterator declIter = decls_.getDeclIter();
        while (declIter.hasNext()) {
            Decl decl = (Decl)declIter.next();
            // if the declaration has already been desugared, no 
            // need to do it again
            if (decl.desugared) continue;
        
            
            MultiplicityExpr multExpr = decl.getExpr();
            Set varExprs = new TreeSet();
            Iterator varIter = decl.getVariables().getVariableIter();
            while (varIter.hasNext()) {
                Variable var = (Variable)varIter.next();
                VariableExpr varExpr = new VariableExpr(var);
                varExpr.setType(multExpr.getType());
                varExprs.add(varExpr);
                Formula constraint;
                
                if (decl.isHigherOrder()) {
                    constraint = new MultiplicityFormula((Expr)varExpr.copy(),
							     (MultiplicityExpr)multExpr.copy());
                
                    constraint = (Formula) constraint.applyReturnVisitor(new DesugarMultFormulasVisitor());
                } else {
                    constraint = new ElemFormula((Expr) varExpr.copy(),
							CompOp.SUBSETCOLON,
							(Expr) multExpr.toExpr().copy());
                }
                
                // newly created decls may need to be desugared
                constraint.applyVisitor(new DesugarDeclsVisitor());
                constraint.annotate("constraints for variable " + var.nodeString());
                allConstraints.addFormula(constraint);
            }
            // if decl was of the form x1...xn : some S replace it with
            // x1...xn : set S, since 'some' was desugared
            if (multExpr instanceof SetMultExpr) {
        		if (multExpr.getMult().equals(Multiplicity.SOME))
        		    multExpr.setMult(Multiplicity.SET);
        	    } else {
        	    // remove all multiplicites from decls of the form x1...xn : m A p->q B so that
        	    // they become x1...xn : A->B
        		ArrowMultExpr arrowExpr = (ArrowMultExpr) multExpr;
        		arrowExpr.setMult(Multiplicity.NONE);
        		arrowExpr.getLeft().setMult(Multiplicity.NONE);
        		arrowExpr.getRight().setMult(Multiplicity.NONE);
        	    }
            if (decl.isDisj() && varExprs.size() > 1) {
                allConstraints.addFormula(ImplicitConstraintGenerator.genDisjConstraints(varExprs));
            }
            if (decl.isExh() && varExprs.size() > 0) {
                allConstraints.addFormula(ImplicitConstraintGenerator.genExhConstraints(varExprs, multExpr.toExpr()));
            }
            decl.desugared = true;
        }
        return new FormulaSeq(allConstraints);
    }
                
}
